{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install langchain_google_vertexai langgraph langchain_chroma"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agentic RAG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a `VectorStore` with some data to illustrate the concept.\n",
    "\n",
    "For that we will use the file `data.json` that contains some descriptions of clothing items and create a list of `Document` objects with it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---\n",
      "Sunny Day Tee: This vibrant yellow t-shirt is perfect for sunny summer days. Made from soft, breathable cotton, it features a classic crew neck and relaxed fit. The cheerful color will brighten up any outfit.\n",
      "---\n",
      "---\n",
      "Denim Dream Pants: These classic blue denim pants are a versatile addition to your fall wardrobe. With a straight leg and mid-rise waist, they offer a comfortable and flattering fit. The durable denim fabric ensures they'll last through many seasons.\n",
      "---\n",
      "---\n",
      "Midnight Magic Coat: Stay warm and stylish in this elegant black winter coat. Crafted from a luxurious wool blend, it features a sleek silhouette and cozy collar. The timeless design makes it a perfect choice for any occasion.\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "from langchain_core.documents import Document\n",
    "\n",
    "with open(\"data.json\") as input_file:\n",
    "    items = json.load(input_file)\n",
    "\n",
    "documents = [Document(f\"{item['name']}: {item['description']}\") for item in items]\n",
    "\n",
    "for document in documents[:3]:\n",
    "    print(\"---\")\n",
    "    print(document.page_content)\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, for illustration porpouses, let's create an in memory `VectorStore` using `Chroma`. Remember that in `Chapter 4` we explored all Vectorstores available in Google Cloud for production applications that follow the same interface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_google_vertexai import VertexAIEmbeddings\n",
    "from langchain.tools.retriever import create_retriever_tool\n",
    "\n",
    "\n",
    "embedding_model = VertexAIEmbeddings(model_name=\"textembedding-gecko@003\")\n",
    "vector_store = Chroma.from_documents(documents=documents, embedding=embedding_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the `VectorStore`, we create a `Tool` that is able to be called by an agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Midnight Magic Coat: Stay warm and stylish in this elegant black winter coat. Crafted from a luxurious wool blend, it features a sleek silhouette and cozy collar. The timeless design makes it a perfect choice for any occasion.\n",
      "\n",
      "Winter Wonderland Jacket: Stay warm and cozy in this winter white jacket. Crafted from a water-resistant fabric, it features a plush lining and faux fur trim. The quilted design adds a touch of style while providing extra insulation against the cold.\n",
      "\n",
      "Cozy Cabin Cardigan: Snuggle up by the fire in this cozy brown cardigan. Knitted from a soft wool blend, it features a chunky knit pattern and oversized fit. The warm neutral color complements a variety of winter outfits.\n",
      "\n",
      "Snowy Peaks Parka: Brave the elements in style with this blue parka. Crafted from a water-resistant and windproof fabric, it features a warm down filling and faux fur-lined hood. The multiple pockets and adjustable cuffs provide added functionality.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "retriever = vector_store.as_retriever(k = 3)\n",
    "\n",
    "retriever_tool = create_retriever_tool(\n",
    "    retriever=retriever,\n",
    "    name=\"retrieve_items\",\n",
    "    description=\"Retrieves clothing item descriptions from the knowledge base.\"\n",
    ")\n",
    "\n",
    "response = retriever_tool.invoke(\"I want a black coat for winter\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now create the agent architecture as we described in the Chapter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "from langchain_google_vertexai import ChatVertexAI\n",
    "from langchain_core.messages import AnyMessage\n",
    "from langgraph.graph import StateGraph, MessagesState, END\n",
    "from langgraph.prebuilt import ToolNode\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "\n",
    "graph = StateGraph(MessagesState)\n",
    "tools = [retriever_tool]\n",
    "\n",
    "generator = ChatVertexAI(\n",
    "  model_name=\"gemini-1.5-pro-001\", \n",
    "  temperature = 0\n",
    ").bind(tools=tools)\n",
    "\n",
    "def invoke_generator(state: MessagesState) -> None:\n",
    "    \"\"\" Represents the generator node.\n",
    "    \"\"\"\n",
    "    response = generator.invoke(state[\"messages\"])\n",
    "    state[\"messages\"].append(response)\n",
    "    \n",
    "def use_retrieval(state: MessagesState) -> Literal[\"retriever_tool\", END]:\n",
    "    \"\"\" Represents the use_retrieval conditional edge\n",
    "    \"\"\"\n",
    "    if not state[\"messages\"]:\n",
    "        return END\n",
    "    if state[\"messages\"][-1].tool_calls:\n",
    "        return \"retriever_tool\"\n",
    "    return END\n",
    "\n",
    "\n",
    "graph.add_node(\"generator\", invoke_generator)\n",
    "graph.add_node(\"retriever_tool\", ToolNode(tools))\n",
    "graph.set_entry_point(\"generator\")\n",
    "\n",
    "graph.add_conditional_edges(\"generator\", use_retrieval)\n",
    "graph.add_edge(\"retriever_tool\", \"generator\")\n",
    "\n",
    "\n",
    "agentic_rag = graph.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can invoke the agent with the `invoke` or `stream` methods. Let's try with a question that doesn't require a function invocation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Hello! How can you help me?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "I can help you navigate a catalog of clothing items. For example, you can ask me to find items similar to \"a red dress\" or \"a pair of blue jeans\". \n",
      "\n",
      "What are you looking for today?\n"
     ]
    }
   ],
   "source": [
    "\n",
    "messages = [\n",
    "    SystemMessage(\n",
    "        \"\"\"\n",
    "        - You are a useful assistant that help users navigate a catalog of clothing items.\n",
    "        - You can retrieve clothing items from the catalog using a natural language query.\n",
    "        - Answer in natural language and format your output using paragraph or bullet points if\n",
    "          necessary.\n",
    "        \"\"\"\n",
    "    ),\n",
    "    HumanMessage(\"Hello! How can you help me?\")\n",
    "]\n",
    "\n",
    "stream_generator = agentic_rag.stream({\"messages\": messages}, stream_mode=\"values\")\n",
    "\n",
    "for state in stream_generator:\n",
    "    state[\"messages\"][-1].pretty_print()\n",
    "state[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we invoke it with a question, it will automatically call the retriever tool:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Do you have any yellow t-shirts?\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: retrieve_items\n",
      "\n",
      "Sunny Day Tee: This vibrant yellow t-shirt is perfect for sunny summer days. Made from soft, breathable cotton, it features a classic crew neck and relaxed fit. The cheerful color will brighten up any outfit.\n",
      "\n",
      "Summer Solstice Tank: Stay cool and comfortable in this classic white tank top. Made from a soft cotton jersey, it features a relaxed fit and scoop neck. The versatile design can be dressed up or down for any occasion.\n",
      "\n",
      "Blossom Blouse: Embrace the beauty of spring with this delicate pink t-shirt. Adorned with a subtle floral print, it exudes femininity and charm. The lightweight fabric and flowy design make it a comfortable and stylish choice.\n",
      "\n",
      "Nautical Vibes Top: Channel your inner sailor with this nautical-inspired blue and white striped t-shirt. Made from a lightweight cotton blend, it's perfect for summer days spent by the sea. The classic design is both timeless and stylish.\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "We have the Sunny Day Tee. It's a vibrant yellow t-shirt perfect for summer, made from soft and breathable cotton with a classic crew neck and relaxed fit.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "messages = [\n",
    "    SystemMessage(\n",
    "        \"\"\"\n",
    "        - You are a useful assistant that help users navigate a catalog of clothing items.\n",
    "        - You can retrieve clothing items from the catalog using a natural language query.\n",
    "        - Answer in natural language and format your output using paragraph or bullet points if\n",
    "          necessary.\n",
    "        \"\"\"\n",
    "    ),\n",
    "    HumanMessage(\"Do you have any yellow t-shirts?\")\n",
    "]\n",
    "\n",
    "stream_generator = agentic_rag.stream({\"messages\": messages}, stream_mode=\"values\")\n",
    "\n",
    "for state in stream_generator:\n",
    "    state[\"messages\"][-1].pretty_print()\n",
    "state[\"messages\"][-1].pretty_print()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
